package com.codezj;

import javassist.*;
import lombok.extern.slf4j.Slf4j;

import java.io.ByteArrayInputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

/**
 * @Author: baozhongjie
 * @Version: v1.0.0
 * @Description: 类转换器
 **/
@Slf4j
public class MethodTransformer implements ClassFileTransformer {


    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        // 代理类或者编译器生成的类我们可以跳过
        if (shouldSkipClass(className)) {
            return null;
        }
        log.info("[INFO] Processing class: {}", className);

        try {
            // 创建独立的 ClassPool，避免缓存问题
            ClassPool pool = new ClassPool(true);

            // 添加应用的类路径，确保能找到所有依赖类
            if (loader != null) {
                pool.appendClassPath(new javassist.LoaderClassPath(loader));
            }

            // 添加系统类路径
            pool.appendSystemPath();

            // 将斜杠格式转换为点格式
            String dotClassName = className.replace('/', '.');

            // 直接从字节数组创建 CtClass
            CtClass ctClass = pool.makeClass(new ByteArrayInputStream(classfileBuffer));

            // 验证类名是否匹配
            if (!ctClass.getName().equals(dotClassName)) {
                log.warn("[WARN] Class name mismatch: expected {}, got {}", dotClassName, ctClass.getName());
                ctClass.detach();
                return null;
            }

            // 跳过接口、注解、枚举和抽象类
            if (ctClass.isInterface() || ctClass.isAnnotation() || ctClass.isEnum()) {
                log.info("[INFO] Skipping interface/annotation/enum: {}", className);
                ctClass.detach();
                return null;
            }

            boolean modified = false;
            CtMethod[] methods = ctClass.getDeclaredMethods();

            log.info("[INFO] Found {} methods in class: {}", methods.length, className);

            for (CtMethod method : methods) {
                // 跳过特殊方法
                if (shouldSkipMethod(method)) {
                    log.info("[INFO] Skipping method: {} in class: {}", method.getName(), className);
                    continue;
                }

                try {
                    enhanceMethod(method);
                    modified = true;
                    log.info("[INFO] Enhanced method: {} in class: {}", method.getName(), className);
                } catch (Exception e) {
                    log.warn("[WARN] Failed to enhance method: {} in class: {}, error: {}",
                            method.getName(), className, e.getMessage());
                }
            }

            if (modified) {
                byte[] bytecode = ctClass.toBytecode();
                ctClass.detach();
                return bytecode;
            } else {
                log.info("[INFO] No methods enhanced in class: {} (found {} methods)", className, methods.length);
                ctClass.detach();
            }

        } catch (Exception e) {
            log.error("[ERROR] Failed to transform class: {}, error: {}: {}",
                    className, e.getClass().getName(), e.getMessage(), e);
        }

        return null;
    }

    /**
     * 增强方法
     * @param method 待增强的CtMethod对象
     */
    private void enhanceMethod(CtMethod method) throws CannotCompileException {
        try {
            // 获取方法信息
            String className = method.getDeclaringClass().getName();
            String methodName = method.getName();

            log.info("[INFO] Enhancing method: {}.{}", className, methodName);

            // 获取方法起始行号
            int startLine = method.getMethodInfo().getLineNumber(0);
            if (startLine <= 0) {
                startLine = method.getMethodInfo().getLineNumber(1);
            }

            // 获取方法参数类型信息
            String paramTypesCode = buildParamTypesCode(method);

            // 方法入口代码 - 使用带参数类型信息的重载方法
            String beforeCode = "{" +
                                "Object[] params = ($args != null) ? $args : new Object[0];" +
                                "Class[] paramTypes = " + paramTypesCode + ";" +
                                "com.codezj.MethodTracer.logMethodEntry(\"" + className + "\", \"" + methodName + "\", params, paramTypes, " + startLine + ");" +
                                "}";

            // 方法正常退出代码
            String afterCode = "{" +
                               "com.codezj.MethodTracer.logMethodExit(\"" + className + "\", \"" + methodName + "\", ($w) $_);" +
                               "}";

            // 方法异常退出代码
            String catchCode = "{" +
                               "com.codezj.MethodTracer.logMethodException(\"" + className + "\", \"" + methodName + "\", $e);" +
                               "throw $e;" +
                               "}";

            // 注入代码
            method.insertBefore(beforeCode);
            method.insertAfter(afterCode);

            // 添加异常处理，捕获所有异常
            method.addCatch(catchCode, method.getDeclaringClass().getClassPool().get("java.lang.Throwable"));

            log.info("[INFO] Successfully enhanced method: {}.{} (line: {}) with complete lifecycle",
                    className, methodName, startLine);

        } catch (CannotCompileException e) {
            log.error("[ERROR] Cannot compile enhanced code for method: {}, error: {}", method.getName(), e.getMessage(), e);
            throw e;
        } catch (Exception e) {
            log.error("[ERROR] Unexpected error enhancing method: {}, error: {}", method.getName(), e.getMessage(), e);
            throw new CannotCompileException(e);
        }
    }

    /**
     * 检查是否为代理类或编译器生成的类
     * @param className 类名
     */
    private boolean isProxyOrGeneratedClass(String className) {
        return className.contains("$EnhancerBySpringCGLIB$") ||
               className.contains("$FastClassBySpringCGLIB$") ||
               className.contains("$Proxy") ||
               className.contains("$Lambda$") ||
               className.contains("$");
    }

    /**
     * 跳过大部分系统类和框架类
     * @param className 类名
     */
    private boolean shouldSkipClass(String className) {
        // 跳过JDK核心类
        if (className.startsWith("java/") ||
            className.startsWith("javax/") ||
            className.startsWith("sun/") ||
            className.startsWith("com/sun/") ||
            className.startsWith("jdk/") ||
            className.startsWith("com/fasterxml")) {
            return true;
        }

        // 跳过常见框架类
        if (className.startsWith("org/springframework/") ||
            className.startsWith("org/apache/") ||
            className.startsWith("org/slf4j/") ||
            className.startsWith("ch/qos/logback/") ||
            className.startsWith("com/fasterxml/jackson/") ||
            className.startsWith("org/hibernate/") ||
            className.startsWith("org/mybatis/") ||
            className.startsWith("com/mysql/") ||
            className.startsWith("org/eclipse/") ||
            className.startsWith("org/junit/") ||
            className.startsWith("org/mockito/") ||
            className.startsWith("org/")) {
            return true;
        }

        // 跳过Agent自身的类
        if (className.startsWith("com/codezj/")) {
            return true;
        }

        // 跳过类加载器相关类
        if (className.startsWith("org/springframework/boot/loader/")) {
            return true;
        }

        if (isProxyOrGeneratedClass(className)) {
            return true;
        }

        // 处理其他类
        return false;
    }


    /**
     * 跳过方法
     * @param method 方法
     */
    private boolean shouldSkipMethod(CtMethod method) {
        try {
            int modifiers = method.getModifiers();
            String methodName = method.getName();

            // 跳过抽象方法、native方法、构造方法、静态初始化方法
            if (Modifier.isAbstract(modifiers) ||
                Modifier.isNative(modifiers) ||
                methodName.equals("<init>") ||
                methodName.equals("<clinit>")) {
                return true;
            }

            // 检查合成方法标志位 (ACC_SYNTHETIC = 0x1000)
            if ((modifiers & 0x1000) != 0) {
                return true;
            }

            // 检查桥接方法标志位 (ACC_BRIDGE = 0x0040)
            if ((modifiers & 0x0040) != 0) {
                return true;
            }

            // 跳过 Lambda 表达式和编译器生成的方法
            if (methodName.startsWith("lambda$") || methodName.startsWith("access$")) {
                return true;
            }

            if (methodName.startsWith("get") || methodName.startsWith("set")) {
                return true;
            }

            return false;
        } catch (Exception e) {
            return true;
        }
    }

    /**
     * 构建参数类型数组的代码字符串
     * @param method 方法对象
     */
    private String buildParamTypesCode(CtMethod method) throws NotFoundException {
        CtClass[] paramTypes = method.getParameterTypes();
        if (paramTypes == null || paramTypes.length == 0) {
            return "new Class[0]";
        }

        StringBuilder sb = new StringBuilder();
        sb.append("new Class[]{");

        for (int i = 0; i < paramTypes.length; i++) {
            if (i > 0) {
                sb.append(",");
            }

            String typeName = paramTypes[i].getName();

            // 处理基本类型
            if (paramTypes[i].isPrimitive()) {
                sb.append(getPrimitiveClassCode(typeName));
            } else {
                // 处理对象类型和数组类型
                sb.append(typeName).append(".class");
            }
        }

        sb.append("}");
        return sb.toString();
    }

    /**
     * 获取基本类型的Class代码
     * @param primitiveTypeName 基本类型名称
     */
    private String getPrimitiveClassCode(String primitiveTypeName) {
        switch (primitiveTypeName) {
            case "boolean":
                return "boolean.class";
            case "byte":
                return "byte.class";
            case "char":
                return "char.class";
            case "double":
                return "double.class";
            case "float":
                return "float.class";
            case "int":
                return "int.class";
            case "long":
                return "long.class";
            case "short":
                return "short.class";
            case "void":
                return "void.class";
        }
        return primitiveTypeName + ".class";
    }
}
